#include <arpa/inet.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <errno.h>
#include <unistd.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#define DEFAULT_HOST "127.0.0.1"
#define DEFAULT_PORT 0x2304

#define CHUNK_COMMAND 0
#define CHUNK_STDIN 1
#define CHUNK_STDOUT 2
#define CHUNK_STDERR 3
#define CHUNK_EXIT 4

#ifndef MIN
#define MIN(a,b) ((a<b)?(a):(b))
#endif

void Writen(int fd, void *vptr, size_t n);
ssize_t Readn(int fd, void *vptr, size_t n);
ssize_t Read(int fd, void *ptr, size_t nbytes);
int Socket(int family, int type, int protocol);
void Inet_pton(int family, const char *strptr, void *addrptr);
void Connect(int fd, const struct sockaddr *sa, socklen_t salen);
int Select(int nfds, fd_set *readfds, fd_set *writefds, fd_set *exceptfds, struct timeval *timeout);

int write_cmd(int fd, char **argv, int argc);
void cmd_loop(int sockfd);

#define BUFLEN 1024

int bytes2short(char *b) {
	uint16_t d = *((uint16_t*)b);
	d = ntohs(d);
	return d;
}

void short2bytes(int d, char *b) {
	uint16_t d1 = htons(d);
	char *data = (char*)&d1;
	b[0] = data[0];
	b[1] = data[1];
}

void print_usage() {
	fprintf(stderr,
		"Usage: sjf [options] command [args]\n"
		"\n"
		"Options:\n"
		"  --help                show this help message and exit\n"
		"  -h HOST               specify the server host to connect\n"
		"  -p PORT               specify the server port to connect, default 0x2304\n"
	);
}

int main(int argc, char **argv) {
	struct sockaddr_in servaddr;
	char *server = DEFAULT_HOST;
	int port = DEFAULT_PORT;
	int sockfd;

	argc--; argv++;
	for (; argc > 0; argc--, argv++) {
		if (strcmp("--help", *argv) == 0) {
			print_usage();
			exit(1);
		} else if (strcmp("-h", *argv) == 0) {
			if (argc > 1) {
				server = *++argv;
				argc--;
			} else {
				print_usage();
				exit(1);
			}
		} else if (strcmp("-p", *argv) == 0) {
			if (argc > 1) {
				port = atoi(*++argv);
				argc--;
			} else {
				print_usage();
				exit(1);
			}
		} else if (*argv == NULL) {
			continue;
		} else if (*argv[0] == '-') {
			print_usage();
			exit(1);
		} else {
			break;
		}
	}

	if (argc < 1) {
		fprintf(stderr, "please specify command\n");
		exit(1);
	}

	sockfd = Socket(AF_INET, SOCK_STREAM, 0);
	memset(&servaddr, '\0', sizeof(servaddr));
	servaddr.sin_family = AF_INET;
	servaddr.sin_port = htons(port);
	Inet_pton(AF_INET, server, &servaddr.sin_addr);
	Connect(sockfd, (struct sockaddr *)&servaddr, sizeof(servaddr));

	write_cmd(sockfd, argv, argc);
	cmd_loop(sockfd);
	exit(0);
}

void write_chunk_header(int fd, int chunk_type, int len) {
	char buf[3];
	buf[0] = chunk_type;
	short2bytes(len, buf+1);
	Writen(fd, buf, 3);
}

void write_short(int fd, int len) {
	char buf[2];
	short2bytes(len, buf);
	Writen(fd, buf, 2);
}

void write_text(int fd, char *str) {
	int len = strlen(str);
	write_short(fd, len);
	Writen(fd, str, len);
}

int write_cmd(int fd, char **argv, int argc) {
	int len;
	char *cwd;

	len = strlen(*argv);
	write_chunk_header(fd, CHUNK_COMMAND, len);
	Writen(fd, *argv++, len);	// write cmd name

	// write args
	write_short(fd, argc-1);
	while (*argv) {
		write_text(fd, *argv++);
	}
	cwd = getcwd(NULL, 0);
	write_text(fd, cwd);
	free(cwd);
	return 0;
}

void cmd_loop(int sockfd) {
	int n, r;
	int chunk_type;
	char buf[BUFLEN];
	fd_set rset;
	
	FD_ZERO(&rset);
	while (1) {
		FD_SET(STDIN_FILENO, &rset);
		FD_SET(sockfd, &rset);
		Select(sockfd+1, &rset, NULL, NULL, NULL);
		if (FD_ISSET(STDIN_FILENO, &rset)) {
			n = Read(STDIN_FILENO, buf, BUFLEN);
			if (n > 0) {
				write_chunk_header(sockfd, CHUNK_STDIN, n);
				Writen(sockfd, buf, n);
			}
		}
		if (FD_ISSET(sockfd, &rset)) {
			if (Readn(sockfd, buf, 3) == 0)		// eof, server closed connection
				return;
			chunk_type = buf[0];
			r = bytes2short(buf+1);
			while (r > 0) {
				int m = MIN(r, BUFLEN);
				n = Readn(sockfd, buf, m);
				if (n == 0) 		// read eof
					return;
				switch(chunk_type) {
				case CHUNK_STDOUT:
					Writen(STDOUT_FILENO, buf, n);
					break;
				case CHUNK_STDERR:
					Writen(STDERR_FILENO, buf, n);
					break;
				case CHUNK_EXIT:
					exit(buf[0]);
				}
				r -= n;
			}
		}
	}
}

// copy from unix network programming
#include	<stdarg.h>		/* ANSI C header file */
#include	<syslog.h>		/* for syslog() */
static void	err_doit(int, int, const char *, va_list);

void err_sys(const char *fmt, ...) {
	va_list		ap;

	va_start(ap, fmt);
	err_doit(1, LOG_ERR, fmt, ap);
	va_end(ap);
	exit(1);
}

void err_quit(const char *fmt, ...) {
	va_list		ap;

	va_start(ap, fmt);
	err_doit(0, LOG_ERR, fmt, ap);
	va_end(ap);
	exit(1);
}

static void err_doit(int errnoflag, int level, const char *fmt, va_list ap) {
	int		errno_save, n;
	char	buf[BUFLEN];

	errno_save = errno;		/* value caller might want printed */
	vsnprintf(buf, sizeof(buf), fmt, ap);	/* this is safe */
	n = strlen(buf);
	if (errnoflag)
		snprintf(buf+n, sizeof(buf)-n, ": %s", strerror(errno_save));
	strcat(buf, "\n");

	fflush(stdout);		/* in case stdout and stderr are the same */
	fputs(buf, stderr);
	fflush(stderr);
	return;
}

static ssize_t writen(int fd, const void *vptr, size_t n) {
	size_t		nleft;
	ssize_t		nwritten;
	const char	*ptr;

	ptr = vptr;
	nleft = n;
	while (nleft > 0) {
		if ( (nwritten = write(fd, ptr, nleft)) <= 0) {
			if (errno == EINTR)
				nwritten = 0;		/* and call write() again */
			else
				return(-1);			/* error */
		}

		nleft -= nwritten;
		ptr   += nwritten;
	}
	return n;
}

void Writen(int fd, void *ptr, size_t nbytes) {
	if (writen(fd, ptr, nbytes) != nbytes)
		err_sys("writen error");
}

ssize_t readn(int fd, void *vptr, size_t n) {
	size_t	nleft;
	ssize_t	nread;
	char	*ptr;

	ptr = vptr;
	nleft = n;
	while (nleft > 0) {
		if ( (nread = read(fd, ptr, nleft)) < 0) {
			if (errno == EINTR)
				nread = 0;		/* and call read() again */
			else
				return(-1);
		} else if (nread == 0)
			break;				/* EOF */

		nleft -= nread;
		ptr   += nread;
	}
	return(n - nleft);		/* return >= 0 */
}

ssize_t Readn(int fd, void *ptr, size_t nbytes) {
	ssize_t		n;

	if ( (n = readn(fd, ptr, nbytes)) < 0)
		err_sys("readn error");
	return(n);
}

ssize_t Read(int fd, void *ptr, size_t nbytes) {
	ssize_t		n;

	if ( (n = read(fd, ptr, nbytes)) == -1)
		err_sys("read error");
	return(n);
}

int Socket(int family, int type, int protocol) {
	int		n;

	if ( (n = socket(family, type, protocol)) < 0)
		err_sys("socket error");
	return(n);
}

void Inet_pton(int family, const char *strptr, void *addrptr) {
	int		n;

	if ( (n = inet_pton(family, strptr, addrptr)) < 0)
		err_sys("inet_pton error for %s", strptr);	/* errno set */
	else if (n == 0)
		err_quit("inet_pton error for %s", strptr);	/* errno not set */

	/* nothing to return */
}

void Connect(int fd, const struct sockaddr *sa, socklen_t salen) {
	if (connect(fd, sa, salen) < 0)
		err_sys("connect error");
}

int Select(int nfds, fd_set *readfds, fd_set *writefds, fd_set *exceptfds,
       struct timeval *timeout) {
	int		n;

	if ( (n = select(nfds, readfds, writefds, exceptfds, timeout)) < 0)
		err_sys("select error");
	return(n);		/* can return 0 on timeout */
}
